import streamlit as st
import pickle
import os
import time
import json
from streamlit import session_state
from zhipuai import ZhipuAI

os.environ['http_proxy'] = 'http://127.0.0.1:7890'
os.environ["https_proxy"] = "http://127.0.0.1:7890"
# streamlit run dataset_maker.py --server.port 2323
st.set_page_config(
    page_title='问答数据集生成器',
    layout="wide",
    initial_sidebar_state="expanded",
    menu_items={
        'Get Help': None,
        'Report a bug': None,
        'About': None
    }
)
# 加载问题库
def load_questions(file_name):
    file_path='./data/question_sets'+file_name
    if not os.path.exists(file_path):
        st.error(f"文件 {file_path} 不存在")
        print(f"文件 {file_path} 不存在")
        return []
    else:
        with open(file_path, "r", encoding='utf-8') as file:
            questions = file.readlines()
        return list(set([q.strip() for q in questions if q != '' and q != '\n']))  # 去重

# 保存问题库
def save_questions(file_path, questions):
    with open(file_path, "w", encoding='utf-8') as file:
        for question in questions:
            file.write(question + "\n")

def save_answers(temp_answers, just_read=False):
    if just_read:
        if os.path.exists("../data/finetune_sets/data.json"):
            with open("../data/finetune_sets/data.json", "rb") as file:
                answers = pickle.load(file)
        else:
            answers = {}
        session_state.all_answers = answers
        return True
    else:
        if os.path.exists("lock"):
            return False
        with open("lock", "w") as lock_file:
            lock_file.write("")
        if os.path.exists("../data/finetune_sets/data.json"):
            with open("../data/finetune_sets/data.json", "rb") as file:
                answers = pickle.load(file)
        else:
            answers = {}
        answers.update(temp_answers)  # 覆盖式更新

        with open("../data/finetune_sets/data.json", "wb") as file:
            pickle.dump(answers, file)
        if os.path.exists("lock"):
            os.remove("lock")
        session_state.all_answers = answers
        return True

def save_answers_as_json(answers, file_path):
    with open(file_path, "w", encoding="utf-8") as file:
        for question, answer in answers.items():
            item = {"prompt": "问题：" + question + "\n", "completion": answer}
            file.write(json.dumps(item, ensure_ascii=False) + "\n")

def reset_text_area():
    if session_state.text_area_tittle == "回答：(内容为空则不保存此回答)":
        session_state.text_area_tittle = "回答：(内容为空则不保存此回答) "
    elif session_state.text_area_tittle == "回答：(内容为空则不保存此回答) ":
        session_state.text_area_tittle = "回答：(内容为空则不保存此回答)"


def main():
    st.title("问答数据集生成器")
    st.markdown("""
        <style>
            .stButton>button {
                background-color: #d4702d; 
                color: white; /* 白色文字 */
                border: none; /* 无边框 */
                padding: 10px 20px; /* 内边距 */
                text-align: center; /* 中心对齐 */
                text-decoration: none; /* 无下划线 */
                display: inline-block; /* 内联块 */
                font-size: 16px; /* 字体大小 */
                margin-left: 1px; /* 外边距 */
                cursor: pointer; /* 鼠标悬停时显示手指 */
                border-radius: 8px; /* 圆角 */
                transition: background-color 0.3s; /* 平滑过渡 */
            }
            .stButton>button:hover {
                background-color: #87481d; /* 悬停时变为更深的绿色 */
                color: #a5aaa3; /* 白色文字 */
            }
        </style>
        """, unsafe_allow_html=True)
    client = ZhipuAI(api_key=st.sidebar.text_input("api_key", value='d1cc37827b72a95d2dba239636037d98.moF6orGZ0KPtm7JD',
                                                   type="password"))
    PROMPT = st.sidebar.text_input("提示词", value="请给出以下问题的答案：")

    if 'temp_answers' not in session_state:
        session_state.temp_answers = {}
    if 'all_answers' not in session_state:
        save_answers(session_state.temp_answers, just_read=True)
        # session_state.question_txt = "questions.txt"
        session_state.answers_json = "answers.json"
        session_state.generated_answer = ""
        session_state.text_area_tittle = "回答：(内容为空则不保存此回答)"

    # session_state.question_txt = st.sidebar.text_input("选择问题集：", value=session_state.question_txt)
    session_state.answers_json = st.sidebar.text_input("保存回答的json路径", value=session_state.answers_json)

    # 选择问题集
    question_files = [f for f in os.listdir("questions") if f.endswith(".txt")]
    selected_file = st.sidebar.selectbox("选择问题集：", question_files)
    session_state.question_txt = selected_file

    if 'questions' not in st.session_state or st.session_state.question_txt != os.path.join("questions", selected_file):
        st.session_state.questions = load_questions(st.session_state.question_txt)

    # 解析问题和序号
    question_options = [(line.split(',', 1)[0], line.split(',', 1)[1].strip()) for line in session_state.questions]

    # 创建选择和输入框
    selected_question_id = st.selectbox("选择或输入您的问题：",
                                        options=["手动输入其他问题"] + [f"{num}. {q}" for num, q in question_options],
                                        format_func=lambda x: x.split(". ", 1)[1] if ". " in x else x)

    # 如果选择了“手动输入其他问题”，则显示输入框
    if selected_question_id == "手动输入其他问题":
        user_input = st.text_input("请手动输入问题：")
    else:
        user_input = selected_question_id.split(". ", 1)[1]

    prompt = PROMPT + user_input
    st.sidebar.write({'提问预览': prompt})
    user_answer_empty = st.empty()
    user_answer = user_answer_empty.text_area(session_state.text_area_tittle, session_state.generated_answer,
                                              height=200)

    col1, col2, col3= st.columns(3)
    with col1:
        if st.button("保存全部回答(超过10个自动保存)") or len(session_state.temp_answers) >= 10:
            session_state.selected_id = 0
            for question in session_state.temp_answers:
                session_state.questions.remove(question)
            save_questions(session_state.question_txt, session_state.questions)
            if save_answers(session_state.temp_answers):
                st.success("全部回答已保存。")
                session_state.temp_answers = {}
            else:
                st.error("保存失败，请稍后重试。")
                time.sleep(5)
            st.rerun()

    with col2:
        if st.button("智谱AI生成回答"):
            session_state.generated_answer = ''
            try:
                response = client.chat.completions.create(
                    model="glm-4-flash",
                    messages=[{"role": "user", "content": prompt}]
                )
                session_state.generated_answer = response.choices[0].message.content
                user_answer = user_answer_empty.text_area(session_state.text_area_tittle,
                                                          session_state.generated_answer, height=200)
            except Exception as e:
                st.error("生成失败，请稍后重试。")

    with col3:
        if st.button("确认此回答(自动下一个)"):
            session_state.generated_answer = ''
            if user_answer != '':
                session_state.temp_answers[user_input] = user_answer
            elif user_input in session_state.temp_answers:
                del session_state.temp_answers[user_input]
            reset_text_area()
            user_answer = user_answer_empty.text_area(session_state.text_area_tittle, height=200)
            st.rerun()

    # 其他按钮处理
    # ...

    # if st.sidebar.button("读一下"):
    #     save_answers(session_state.temp_answers, just_read=True)
    #     session_state.questions = load_questions(session_state.question_txt)
    #     st.rerun()

    if st.sidebar.button("导出载入的已保存回答为 JSON"):
        save_answers_as_json(session_state.all_answers, session_state.answers_json)

    st.json({"未保存回答：": session_state.temp_answers, "已保存回答：": session_state.all_answers})


if __name__ == "__main__":
    main()
